package gorpc

import (
	"errors"
	"fmt"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

// RPCClient implements a WebSocket-based RPC client
type RPCClient struct {
	url               string
	headers           map[string]string
	conn              *websocket.Conn
	callback          RPCCallback
	writeMutex        sync.Mutex
	closed            atomic.Bool
	closeMutex        sync.Mutex
	callContexts      map[int32]*RPCCallContext
	callContextsMutex sync.RWMutex
	nextSeqID         atomic.Int32
	session           *RPCSession
}

// NewRPCClient creates a new RPC client instance
func NewRPCClient(callback RPCCallback) *RPCClient {
	if callback == nil {
		callback = NewBaseCallback()
	}
	client := &RPCClient{
		callback:     callback,
		headers:      make(map[string]string),
		closed:       atomic.Bool{},
		callContexts: make(map[int32]*RPCCallContext),
	}
	client.closed.Store(true) // Initialize closed to true
	return client
}

// SetHeader sets a custom header for the WebSocket connection
func (c *RPCClient) SetHeader(key, value string) {
	c.headers[key] = value
}

// Connect establishes a connection to the RPC server
func (c *RPCClient) Connect(url string) error {
	c.closeMutex.Lock()
	if !c.closed.Load() {
		c.closeMutex.Unlock()
		return errors.New("rpc: already connected")
	}
	c.closed.Store(false)
	c.closeMutex.Unlock()

	c.url = url

	// Create HTTP header
	header := http.Header{}
	for k, v := range c.headers {
		header.Add(k, v)
	}

	// Connect to WebSocket server
	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}

	conn, _, err := dialer.Dial(url, header)
	if err != nil {
		c.closed.Store(true)
		return fmt.Errorf("rpc: dial error: %w", err)
	}

	c.conn = conn

	// Create and initialize session
	c.session = NewRPCSession(conn)

	// 设置会话回调
	c.session.SetRequestCallback(func(inMsg *RPCMessage, outMsg *RPCMessage) error {
		if c.callback != nil {
			return c.callback.OnRequest(c.session, inMsg, outMsg, 30*time.Second)
		}
		return nil
	})

	c.session.SetNotifyCallback(func(msg *RPCMessage) {
		if c.callback != nil {
			c.callback.OnNotify(c.session, msg)
		}
	})

	c.session.SetConnectedCallback(func() {
		if c.callback != nil {
			c.callback.OnConnect(c.session)
		}
	})

	c.session.SetDisconnectedCallback(func(err error) {
		if c.callback != nil {
			var code int
			var msg string
			if err != nil {
				code = websocket.CloseAbnormalClosure
				msg = err.Error()
			} else {
				code = websocket.CloseNormalClosure
				msg = "connection closed"
			}
			c.callback.OnClose(c.session, code, msg)
		}
	})

	// 启动会话
	go c.session.Start()

	return nil
}

// Close terminates the RPC connection
func (c *RPCClient) Close(reason string) error {
	c.closeMutex.Lock()
	defer c.closeMutex.Unlock()

	if c.closed.Load() {
		return nil
	}

	c.closed.Store(true)

	if c.session != nil {
		c.session.Close()
	}

	return nil
}

// IsConnected checks if the client is connected
func (c *RPCClient) IsConnected() bool {
	c.closeMutex.Lock()
	defer c.closeMutex.Unlock()
	return !c.closed.Load() && c.conn != nil
}

// WriteMessage sends a WebSocket message
func (c *RPCClient) WriteMessage(data []byte) error {
	if !c.IsConnected() {
		return errors.New("not connected")
	}

	c.writeMutex.Lock()
	defer c.writeMutex.Unlock()

	return c.conn.WriteMessage(websocket.BinaryMessage, data)
}

// CallMethod performs an RPC call with message objects
func (c *RPCClient) CallMethod(inMsg *RPCMessage, outMsg *RPCMessage, timeout time.Duration) error {
	if c.closed.Load() {
		return errors.New("rpc: client closed")
	}

	if c.session == nil {
		return errors.New("rpc: session not established")
	}

	return c.session.CallMethod(inMsg, outMsg, timeout)
}

// Notify sends a notification message
func (c *RPCClient) Notify(msg *RPCMessage) error {
	if c.closed.Load() {
		return errors.New("rpc: client closed")
	}

	if c.session == nil {
		return errors.New("rpc: session not established")
	}

	return c.session.Notify(msg)
}

// AddCallContext adds a call context
func (c *RPCClient) AddCallContext(ctx *RPCCallContext) {
	c.callContextsMutex.Lock()
	defer c.callContextsMutex.Unlock()
	c.callContexts[ctx.SequenceID()] = ctx
}

// RemoveCallContext removes a call context
func (c *RPCClient) RemoveCallContext(seqID int32) {
	c.callContextsMutex.Lock()
	defer c.callContextsMutex.Unlock()
	delete(c.callContexts, seqID)
}

// ClearAllCallContexts clears all call contexts
func (c *RPCClient) ClearAllCallContexts() {
	c.callContextsMutex.Lock()
	defer c.callContextsMutex.Unlock()

	networkErr := fmt.Errorf("%w: connection closed", ErrNetworkDown)

	for _, ctx := range c.callContexts {
		ctx.ForceComplete(networkErr)
	}

	c.callContexts = make(map[int32]*RPCCallContext)
}
